package rabbit.flt.plugins.webflux.plugin;

import rabbit.flt.plugins.common.ContextName;
import rabbit.flt.plugins.common.WebClientCallBack;
import rabbit.flt.plugins.common.plugin.SupportPlugin;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.lang.reflect.Method;

public class WebClientPlugin extends SupportPlugin {

    @Override
    public Object after(Object objectEnhanced, Method method, Object[] args, Object result) {
        super.after(objectEnhanced, method, args, result);
        WebClientCallBack callBack = new WebClientCallBack();
        if (result instanceof Mono) {
            return ((Mono<?>) result).contextWrite(ctx -> ctx.put(ContextName.WEB_CLIENT_CALL_BACK, callBack))
                    .doFinally(signalType -> callBack.getCallBack().run());
        } else {
            return ((Flux<?>) result).contextWrite(ctx -> ctx.put(ContextName.WEB_CLIENT_CALL_BACK, callBack))
                    .doFinally(signalType -> callBack.getCallBack().run());
        }
    }
}
